#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#define CI_TILE 4
#define H 36

__constant sampler_t
smp_none = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;

__kernel void WinogradConvolution(__read_only image2d_t input,
                                  __write_only image2d_t output,
                                  __global FLT16 *weight,
                                  int4 input_shape,         // N 36 TILE_XY CI_SLICES
                                  int4 output_shape)        // N 36 TILE_XY CO_SLICES
{
    int w = get_global_id(0) * 2;
    int h = get_global_id(1);
    int co_slice = get_global_id(2) * 2;

    int CI_SLICES = input_shape.w;
    int W = input_shape.z;
    int CO_SLICES = output_shape.w;

    if (h >= H || w >= W || co_slice >= CO_SLICES)
    {
        return;
    }

    FLT4 out00 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
    FLT4 out01 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
    FLT4 out10 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
    FLT4 out11 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);

    int y_idx = h;
    __global FLT16 *weight_ptr = weight + (co_slice / 2 * 36 + h) * CI_SLICES * 2;
    for (int ci_slice = 0; ci_slice < CI_SLICES; ci_slice++)
    {
        FLT4 in0 = READ_IMAGE(input, smp_none, (int2)(w + 0, y_idx));
        FLT4 in1 = READ_IMAGE(input, smp_none, (int2)(w + 1, y_idx));
        y_idx += 36;

        FLT16 weight0 = weight_ptr[0], weight1 = weight_ptr[1];
        weight_ptr += 2;


        out00 += in0.x * weight0.s0123;
        out00 += in0.y * weight0.s4567;
        out00 += in0.z * weight0.s89ab;
        out00 += in0.w * weight0.scdef;

        out01 += in1.x * weight0.s0123;
        out01 += in1.y * weight0.s4567;
        out01 += in1.z * weight0.s89ab;
        out01 += in1.w * weight0.scdef;

        out10 += in0.x * weight1.s0123;
        out10 += in0.y * weight1.s4567;
        out10 += in0.z * weight1.s89ab;
        out10 += in0.w * weight1.scdef;

        out11 += in1.x * weight1.s0123;
        out11 += in1.y * weight1.s4567;
        out11 += in1.z * weight1.s89ab;
        out11 += in1.w * weight1.scdef;
    }

    WRITE_IMAGE(output, (int2)(w + 0, (co_slice + 0) * H + h), out00);
    if (w + 1 < W)
    {
        WRITE_IMAGE(output, (int2)(w + 1, (co_slice + 0) * H + h), out01);
    }

    if (co_slice + 1 < CO_SLICES)
    {
        WRITE_IMAGE(output, (int2)(w + 0, (co_slice + 1) * H + h), out10);
        if (w + 1 < W)
        {
            WRITE_IMAGE(output, (int2)(w + 1, (co_slice + 1) * H + h), out11);
        }
    }
}